﻿#include "MadelungSystem.h"


#include "Application.h"
#include "Mesh.h"
#include "Transient.h"
#include <vtkm/VectorAnalysis.h>
#include <vtkm/cont/Algorithm.h>
#include <vtkm/worklet/WorkletCellNeighborhood.h>
#include <vtkm/worklet/WorkletMapField.h>
#include <vtkm/worklet/WorkletPointNeighborhood.h>

struct MadelungSystem::InitialCondtionWorklet : vtkm::worklet::WorkletPointNeighborhood
{
  using ControlSignature = void(CellSetIn, FieldOut charge);
  using ExecutionSignature = void(Boundary, _2);

  template<typename T>
  VTKM_EXEC void operator()(const vtkm::exec::BoundaryState& boundary, T& charge) const
  {
    if ((boundary.IJK[0] + boundary.IJK[1] + boundary.IJK[2]) % 2 == 0)
      charge = -1;
    else
    {
      charge = 1;
    }
  }
};

struct MadelungSystem::ComputeCoeffient : vtkm::worklet::WorkletMapField
{
  using ControlSignature = void(FieldIn pts, FieldIn charge, FieldOut force);
  using ExecutionSignature = void(_1, _2, _3);

  template<typename CoordType>
  VTKM_EXEC void operator()(const CoordType& pts, const Real& charge, Real& force) const
  {
    auto p = pts - _pts;
    Real dis = vtkm::Magnitude(p);
    if (dis < 0.0001)
      force = 0;
    else
      force = charge / dis;
  }
  Vec3f _pts{ 0, 0, 0 };
};


RegisterObject(MadelungSystem);

MadelungSystem::MadelungSystem(const Configuration& cfg)
  : System(cfg)
{
}

void MadelungSystem::Init()
{
  System::Init();
  _data_set.SetCellSet(_mesh.GetCellSet());
  _data_set.AddCoordinateSystem(_mesh.GetCoordinate());
  this->AddPointVariable<Real>("charge");

  auto charge = this->GetVariableArray<Real>("charge");
  
  vtkm::cont::Invoker{}(MadelungSystem::InitialCondtionWorklet{}, _data_set.GetCellSet(), charge);
}

void MadelungSystem::Evolve()
{
  _timer.Start();

  PreSolve();
  Solve();
  PostSolve();
  auto coord = _mesh.GetCoordinate().GetDataAsMultiplexer();
  auto charge = this->GetVariableArray<Real>("charge");
  ArrayHandle<Real> force;
  vtkm::cont::Invoker{}(MadelungSystem::ComputeCoeffient{}, coord, charge, force);

  auto result = vtkm::cont::Algorithm::Reduce(force, vtkm::TypeTraits<Real>::ZeroInitialization());
  //std::cout << force.ReadPortal().Get(10) << std::endl;

  std::cout << "Madelung常数：" <<  result << std::endl;

  Info("MadelungSystem::Evolve运行完成，", "计算时间: ", _timer.GetElapsedTime(), " s");
}

void MadelungSystem::PreSolve() {}


void MadelungSystem::TimeIntegration() {}

void MadelungSystem::Solve() {}

void MadelungSystem::PostSolve() {}
